#include "include/types.h"
#include "include/mm.h"

struct bucket_desc { /* Page bucket unit */
  void* freeptr; /* (8B)[VA] Pointer to the next free bucket within the page */
  struct bucket_desc* next; /* (8B)[VA] Pointer to the next page bucket, only the first bucket on the page has this field */
  short refcnt; /* (2B)Record the total occupancy of the bucket, only the first bucket on the page has this field */
};

struct bucket_dir { /* Page bucket byte size unit */
  short size; /* (2B)Page bucket byte size */
  struct bucket_desc* bdesc; /* (8B)Page bucket unit */
};

struct bucket_dir bucket_dir[] = { /* Multiple page bucket byte size arrays */
  { 32, NULL },
  { 64, NULL },
  { 128, NULL },
  { 256, NULL },
  { 512, NULL },
  { 1024, NULL },
  { 2048, NULL },
  { 0, NULL }, /* Traversing termination placeholder unit */
};

/*
 * malloc memory page bucket unit function
 *
 * @param size Memory block byte size to be allocated
 * @return The virtual address of memory block unit
 */
void* malloc(int size) {
  void* va = NULL; /* The virtual address of memory block unit to be returned */
  struct bucket_dir* bdir = NULL; /* Pointer to bucket byte size array */
  struct bucket_desc* bdesc = NULL; /* Pointer to specific sized bucket unit */

  /* Traverse bucket_dir array to find the appropriate bucket unit size */
  for (bdir = bucket_dir; bdir->size; bdir++) {
    if (bdir->size >= size) {
      break;
    }
  }

  if (bdir->size == 0) { /* End if not found */
    return NULL;
  }

  /* Traverse bucket_desc linked list to find page with idle buckets */
  for (bdesc = bdir->bdesc; bdesc; bdesc = bdesc->next) {
    if (bdesc->freeptr) {
      break;
    }
  }

  if (bdesc == NULL) { /* Reassign new page if not found */
    bdesc = VA(alloc_page()); /* (VA) Pointer to newly assigned page */
    /* (Header) Pointer to first free bucket on the page */
    bdesc->freeptr = (void*)((unsigned long)bdesc + sizeof(struct bucket_desc));
    bdesc->next = NULL;
    bdesc->refcnt = 0; /* init size 0 */

    /* Create a linked list of free buckets on the page */
    unsigned long *p = bdesc->freeptr;
    int i = 0; /* bucket index */
    /* ++i: Only traverse to the second to last free bucket */
    /* Remove the first header bucket size before calculating the total number of available buckets */
    while (++i < (PAGE_SIZE - sizeof(struct bucket_desc)) / bdir->size) {
      *p = (unsigned long)p + bdir->size; /* Set the next free bucket virtual address to the first 8B (field freeptr) */
      p = (unsigned long*)(*p); /* Pointer moves to the next free bucket virtual address */
    }
    *p = 0; /* The last free bucket points to NULL(0) */

    /* Insert the new allocation page bdesc into the chain header of bdir */
    bdesc->next = bdir->bdesc;
    bdir->bdesc = bdesc;
  }

  va = bdesc->freeptr; /* Get the first free bucket on the page */
  /* Remove the first one from the idle bucket linked list */
  bdesc->freeptr = (void*)(*((unsigned long*)va)); /* Take the next address value at the address pointed to by va */
  bdesc->refcnt++; /* Bucket occupancy plus 1 */

  return va;
}

/*
 * free memory page bucket unit function
 *
 * @param va The virtual address of memory block unit
 */
void free(void* va) {
  unsigned long page = (unsigned long)va & 0xfffffffffffff000; /* Convert to page base address, 4KB address alignment, using bottom 3 bits of hexadecimal */
  struct bucket_dir* bdir = bucket_dir; /* Pointer to bucket byte size array */
  struct bucket_desc* bdesc = NULL; /* Pointer to specific sized bucket unit */
  /* Traverse bucket_dir to find the corresponding page and return the bucket to it */
  for (; bdir->size; bdir++) {
    for (bdesc = bdir->bdesc; bdesc; bdesc = bdesc->next) {
      if ((unsigned long)bdesc == page) {
        /* Insert the returned bucket into the header of the page bucket */
        *((unsigned long*)va) = (unsigned long)(bdesc->freeptr);
        bdesc->freeptr = va;
        bdesc->refcnt--; /* Bucket occupancy minus 1 */
        break;
      }
    }
  }
  /* Free if the page is idle */
  if (bdesc && bdesc->refcnt == 0) {
    struct bucket_desc* tmp = bdir->bdesc;
    struct bucket_desc* prev = NULL;
    /* Traverse page linked list to find the precursor page of current bdesc */
    for (; tmp; tmp = tmp->next) {
      if ((unsigned long)tmp == (unsigned long)bdesc) {
        break;
      }
      prev = tmp;
    }
    /* Remove the current page(bdesc) from the page linked list */
    if (!prev) { /* not exists precursor */
      bdir->bdesc = bdesc->next;
    } else { /* exists precursor */
      prev->next = bdesc->next;
    }
    /* Free page */
    free_page((unsigned long)bdesc);
  }
}